// For more information, see https://crawlee.dev/
import { Configuration, PlaywrightCrawler, downloadListOfUrls } from "crawlee";
import { readFile, writeFile } from "fs/promises";
import { glob } from "glob";
import { Config, configSchema } from "./config.js";
import { Page } from "playwright";
import { isWithinTokenLimit } from "gpt-tokenizer";
import { PathLike } from "fs";

let pageCounter = 0;
let crawler: PlaywrightCrawler;

export function getPageHtml(page: Page, selector = "body") {
  // evaluate 内部为浏览器上下文直接执行, 无法使用 ts、外部函数
  return page.evaluate(
    ({ selector }) => {
      // @ts-ignore
      function getContentWithAbsoluteLinks(el) {
        // @ts-ignore
        el?.querySelectorAll("a[href]").forEach((a) => {
          const anchorElement = a;
          anchorElement.href = anchorElement.href;
        });
        // 移除 svg iframe 等无用元素
        el
          ?.querySelectorAll("svg,iframe,embed,script,style")
          // @ts-ignore
          .forEach((impurityEl) => {
            impurityEl.remove();
          });

        // 转换所有图片链接 <img> 为绝对链接, 并排除 base64 或超长链接
        // @ts-ignore
        el?.querySelectorAll("img[src]").forEach((img) => {
          const imgElement = img;
          const isBase64 = imgElement.src.startsWith("data:image");
          const isSVG =
            imgElement.src.endsWith(".svg") ||
            imgElement.src.startsWith("data:image/svg+xml");
          const isTooLong = imgElement.src.length > 1024;
          if (isBase64 || isSVG || isTooLong) {
            imgElement.remove();
          } else {
            imgElement.src = imgElement.src; // 转换为绝对链接
          }
        });
        return el;
      }

      // Check if the selector is an XPath
      if (selector.startsWith("/")) {
        const elements = document.evaluate(
          selector,
          document,
          null,
          XPathResult.ANY_TYPE,
          null,
        );
        let result = elements.iterateNext();
        return result ? result.textContent || "" : "";
      } else {
        let el = Array.from(document.querySelectorAll(selector)).slice(-1)[0];
        // 使用 JS 将所有链接（包括图片链接和超链接）转换为绝对链接
        if (el) {
          el = getContentWithAbsoluteLinks(el);
        }
        return el?.innerHTML || "";
      }
    },
    { selector },
  );
}

export async function waitForXPath(page: Page, xpath: string, timeout: number) {
  await page.waitForFunction(
    (xpath) => {
      const elements = document.evaluate(
        xpath,
        document,
        null,
        XPathResult.ANY_TYPE,
        null,
      );
      return elements.iterateNext() !== null;
    },
    xpath,
    { timeout },
  );
}
// 自定义实现，等待异步请求完成。使用内置 networkidle 可能导致导致想要等待异步请求时，页面存在视频、图片持续加载导致抓取网页慢、超时
async function waitForAsyncRequestsToFinish(page: Page, timeout: number = 15000): Promise<void> {
  const pendingRequests = new Set<string>();

  page.on('request', (request) => {
    if (['xhr', 'fetch'].includes(request.resourceType())) {
      pendingRequests.add(request.url());
    }
  });

  page.on('requestfinished', (request) => {
    pendingRequests.delete(request.url());
  });

  page.on('requestfailed', (request) => {
    pendingRequests.delete(request.url());
  });

  const startTime = Date.now();

  // 循环检查 pendingRequests 是否为空
  while (pendingRequests.size > 0) {
    if (Date.now() - startTime > timeout) {
      throw new Error('异步请求超时');
    }
    await page.waitForTimeout(100);  // 每隔100ms检查一次
  }
}

export async function crawl(config: Config) {
  configSchema.parse(config);

  if (process.env.NO_CRAWL !== "true") {
    // PlaywrightCrawler crawls the web using a headless
    // browser controlled by the Playwright library.
    crawler = new PlaywrightCrawler(
      {
        // Use the requestHandler to process each of the crawled pages.
        async requestHandler({ request, page, enqueueLinks, log, pushData }) {
          // 抓取前等待 fetch xhr 完成。
          try {
            // await waitForAsyncRequestsToFinish(page);
            await page.waitForLoadState("networkidle", { timeout: 15000 });
            // page.waitForResponse((response) => { });
          } catch (error) {
            console.log(error);
          }

          // await page.waitForLoadState("networkidle", { timeout: 15000 });
          const title = await page.title();
          pageCounter++;
          log.info(
            `Crawling: Page ${pageCounter} / ${config.maxPagesToCrawl} - URL: ${request.loadedUrl}...`,
          );

          // Use custom handling for XPath selector
          if (config.selector) {
            if (config.selector.startsWith("/")) {
              await waitForXPath(
                page,
                config.selector,
                config.waitForSelectorTimeout ?? 1000,
              );
            } else {
              await page.waitForSelector(config.selector, {
                timeout: config.waitForSelectorTimeout ?? 1000,
              });
            }
          }

          let html = await getPageHtml(page, config.selector);
          await pushData({ 标题: title, 链接: request.loadedUrl, html });

          if (config.onVisitPage) {
            await config.onVisitPage({ page, pushData });
          }

          // Extract links from the current page
          // and add them to the crawling queue.
          await enqueueLinks({
            globs:
              typeof config.match === "string" ? [config.match] : config.match,
            exclude:
              typeof config.exclude === "string"
                ? [config.exclude]
                : config.exclude ?? [],
          });
        },
        // Comment this option to scrape the full website.
        maxRequestsPerCrawl: config.maxPagesToCrawl,
        // Uncomment this option to see the browser window.
        // headless: false,
        preNavigationHooks: [
          // Abort requests for certain resource types
          async ({ request, page, log }) => {
            // If there are no resource exclusions, return
            const RESOURCE_EXCLUSTIONS = config.resourceExclusions ?? [];
            if (RESOURCE_EXCLUSTIONS.length === 0) {
              return;
            }
            if (config.cookie) {
              const cookies = (
                Array.isArray(config.cookie) ? config.cookie : [config.cookie]
              ).map((cookie) => {
                return {
                  name: cookie.name,
                  value: cookie.value,
                  url: request.loadedUrl,
                };
              });
              await page.context().addCookies(cookies);
            }
            await page.route(
              `**\/*.{${RESOURCE_EXCLUSTIONS.join()}}`,
              (route) => route.abort("aborted"),
            );
            log.info(
              `Aborting requests for as this is a resource excluded route`,
            );
          },
        ],
      },
      new Configuration({
        purgeOnStart: true,
      }),
    );

    const isUrlASitemap = /sitemap.*\.xml$/.test(config.url);

    if (isUrlASitemap) {
      const listOfUrls = await downloadListOfUrls({ url: config.url });

      // Add the initial URL to the crawling queue.
      await crawler.addRequests(listOfUrls);

      // Run the crawler
      await crawler.run();
    } else {
      // Add first URL to the queue and start the crawl.
      await crawler.run([config.url]);
    }
  }
}

export async function write(config: Config) {
  let nextFileNameString: PathLike = "";
  const jsonFiles = await glob("storage/datasets/default/*.json", {
    absolute: true,
  });

  console.log(`Found ${jsonFiles.length} files to combine...`);

  let currentResults: Record<string, any>[] = [];
  let currentSize: number = 0;
  let fileCounter: number = 1;
  const maxBytes: number = config.maxFileSize
    ? config.maxFileSize * 1024 * 1024
    : Infinity;

  const getStringByteSize = (str: string): number =>
    Buffer.byteLength(str, "utf-8");

  const nextFileName = (): string =>
    `${config.outputFileName.replace(/\.json$/, "")}-${fileCounter}.json`;

  const writeBatchToFile = async (): Promise<void> => {
    nextFileNameString = nextFileName();
    await writeFile(
      nextFileNameString,
      JSON.stringify(currentResults, null, 2),
    );
    console.log(
      `Wrote ${currentResults.length} items to ${nextFileNameString}`,
    );
    currentResults = [];
    currentSize = 0;
    fileCounter++;
  };

  let estimatedTokens: number = 0;

  const addContentOrSplit = async (
    data: Record<string, any>,
  ): Promise<void> => {
    const contentString: string = JSON.stringify(data);
    const tokenCount: number | false = isWithinTokenLimit(
      contentString,
      config.maxTokens || Infinity,
    );

    if (typeof tokenCount === "number") {
      if (estimatedTokens + tokenCount > config.maxTokens!) {
        // Only write the batch if it's not empty (something to write)
        if (currentResults.length > 0) {
          await writeBatchToFile();
        }
        // Since the addition of a single item exceeded the token limit, halve it.
        estimatedTokens = Math.floor(tokenCount / 2);
        currentResults.push(data);
      } else {
        currentResults.push(data);
        estimatedTokens += tokenCount;
      }
    }

    currentSize += getStringByteSize(contentString);
    if (currentSize > maxBytes) {
      await writeBatchToFile();
    }
  };

  // Iterate over each JSON file and process its contents.
  for (const file of jsonFiles) {
    const fileContent = await readFile(file, "utf-8");
    const data: Record<string, any> = JSON.parse(fileContent);
    await addContentOrSplit(data);
  }

  // Check if any remaining data needs to be written to a file.
  if (currentResults.length > 0) {
    await writeBatchToFile();
  }

  return nextFileNameString;
}

class GPTCrawlerCore {
  config: Config;

  constructor(config: Config) {
    this.config = config;
  }

  async crawl() {
    await crawl(this.config);
  }

  async write(): Promise<PathLike> {
    // we need to wait for the file path as the path can change
    return new Promise((resolve, reject) => {
      write(this.config)
        .then((outputFilePath) => {
          resolve(outputFilePath);
        })
        .catch(reject);
    });
  }
}

export default GPTCrawlerCore;
